import pickle 
import pandas as pd
import numpy as np
from numpy.core.fromnumeric import size
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import train_test_split
from sklearn.metrics import precision_score
from sklearn.metrics import accuracy_score
from sklearn.metrics import f1_score
from sklearn.metrics import confusion_matrix
from sklearn.metrics import recall_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import GaussianNB

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

dtypes = {
    'Src IP': 'category',
    'Src Port': 'uint16',
    'Dst IP': 'category',
    'Dst Port': 'uint16',
    'Protocol': 'category',
    'Flow Duration': 'uint32',
    'Tot Fwd Pkts': 'uint32',
    'Tot Bwd Pkts': 'uint32',
    'TotLen Fwd Pkts': 'float32',
    'TotLen Bwd Pkts': 'float32',
    'Fwd Pkt Len Max': 'float32',
    'Fwd Pkt Len Min': 'float32',
    'Fwd Pkt Len Mean': 'float32',
    'Fwd Pkt Len Std': 'float32',
    'Bwd Pkt Len Max': 'float32',
    'Bwd Pkt Len Min': 'float32',
    'Bwd Pkt Len Mean': 'float32',
    'Bwd Pkt Len Std': 'float32',
    'Flow Byts/s': 'float32',
    'Flow Pkts/s': 'float32',
    'Flow IAT Mean': 'float32',
    'Flow IAT Std': 'float32',
    'Flow IAT Max': 'float32',
    'Flow IAT Min': 'float32',
    'Fwd IAT Tot': 'float32',
    'Fwd IAT Mean': 'float32',
    'Fwd IAT Std': 'float32',
    'Fwd IAT Max': 'float32',
    'Fwd IAT Min': 'float32',
    'Bwd IAT Tot': 'float32',
    'Bwd IAT Mean': 'float32',
    'Bwd IAT Std': 'float32',
    'Bwd IAT Max': 'float32',
    'Bwd IAT Min': 'float32',
    'Fwd PSH Flags': 'category',
    'Bwd PSH Flags': 'category',
    'Fwd URG Flags': 'category',
    'Bwd URG Flags': 'category',
    'Fwd Header Len': 'uint32',
    'Bwd Header Len': 'uint32',
    'Fwd Pkts/s': 'float32',
    'Bwd Pkts/s': 'float32',
    'Pkt Len Min': 'float32',
    'Pkt Len Max': 'float32',
    'Pkt Len Mean': 'float32',
    'Pkt Len Std': 'float32',
    'Pkt Len Var': 'float32',
    'FIN Flag Cnt': 'category',
    'SYN Flag Cnt': 'category',
    'RST Flag Cnt': 'category',
    'PSH Flag Cnt': 'category',
    'ACK Flag Cnt': 'category',
    'URG Flag Cnt': 'category',
    'CWE Flag Count': 'category',
    'ECE Flag Cnt': 'category',
    'Down/Up Ratio': 'float32',
    'Pkt Size Avg': 'float32',
    'Fwd Seg Size Avg': 'float32',
    'Bwd Seg Size Avg': 'float32',
    'Fwd Byts/b Avg': 'uint32',
    'Fwd Pkts/b Avg': 'uint32',
    'Fwd Blk Rate Avg': 'uint32',
    'Bwd Byts/b Avg': 'uint32',
    'Bwd Pkts/b Avg': 'uint32',
    'Bwd Blk Rate Avg': 'uint32',
    'Subflow Fwd Pkts': 'uint32',
    'Subflow Fwd Byts': 'uint32',
    'Subflow Bwd Pkts': 'uint32',
    'Subflow Bwd Byts': 'uint32',
    'Init Fwd Win Byts': 'uint32',
    'Init Bwd Win Byts': 'uint32',
    'Fwd Act Data Pkts': 'uint32',
    'Fwd Seg Size Min': 'uint32',
    'Active Mean': 'float32',
    'Active Std': 'float32',
    'Active Max': 'float32',
    'Active Min': 'float32',
    'Idle Mean': 'float32',
    'Idle Std': 'float32',
    'Idle Max': 'float32',
    'Idle Min': 'float32',
    'Label': 'category'
}

print("Lendo conjunto de dados...")

df = pd.read_csv('./dataset.csv',
     dtype=dtypes,
     parse_dates=['Timestamp'],
     usecols=[*dtypes.keys(), 'Timestamp'],
     engine='c',
     low_memory=True
     )

print("Leitura do CSV via Pandas sucedida!")

def ajusta_dataset(df):
    print("Iniciando procedimentos de ajustes no dataset...")
    # Remoção das colunas com apenas um valor visualizado.
    colsToDrop = np.array(['Fwd Byts/b Avg', 'Fwd Pkts/b Avg', 'Fwd Blk Rate Avg', 'Bwd Byts/b Avg', 'Bwd Pkts/b Avg', 'Bwd Blk Rate Avg'])

    # Remoção de colunas onde os valores ausentes são mais de 40%.
    # Remoção de linhas onde os valores ausentes de uma coluna sejam superiores a 5%.
    missing = df.isna().sum()
    missing = pd.DataFrame({'count': missing, '% of total': missing/len(df)*100}, index=df.columns)
    colsToDrop = np.union1d(colsToDrop, missing[missing['% of total'] >= 40].index.values)
    dropnaCols = missing[(missing['% of total'] > 0) & (missing['% of total'] <= 5)].index.values

    # Ajustando dados 'faltosos'.
    df['Flow Byts/s'].replace(np.inf, np.nan, inplace=True)
    df['Flow Pkts/s'].replace(np.inf, np.nan, inplace=True)
    dropnaCols = np.union1d(dropnaCols, ['Flow Byts/s', 'Flow Pkts/s'])

    df.drop(columns=colsToDrop, inplace=True)
    df.dropna(subset=dropnaCols, inplace=True)
    print("Fim dos ajustes. \n")
    return df

def treina_modelo(df):
    print("Treinamento do modelo:")
    atributos_features_recursos = ["Fwd Seg Size Avg", "Flow IAT Min", "Flow Duration", "Tot Fwd Pkts", "Pkt Size Avg", "Src Port", "Init Bwd Win Byts"]
    alvo_predicao = "Label"

    # Fwd Seg Size Avg - Tamanho médio observado.
    # Flow IAT Min - Tempo mínimo entre dois fluxos de pacotes.
    # Flow Duration - Duração de um fluxo de pacotes.
    # Tot Fwd Pkts - Total de pacotes.
    # Pkt Size Avg - Tamanho médio do pacote.
    # Src Port - Porta de saída no hospedeiro de origem.
    # Label - Status (ddos ou benign) // Valor a ser 'previsto'.

    x = df.loc[:,atributos_features_recursos]
    y = df.loc[:,alvo_predicao]

    from sklearn.tree import DecisionTreeClassifier
    model = DecisionTreeClassifier() # ou GaussianNB ou RandomForestClassifier

    # Treinamento com sklearn, conforme definido.
    print("Treinando modelo...")
    x_train, x_test, y_train, y_test = train_test_split(x, y, random_state=42)   

    model.fit(x_train.values, y_train.values)
    y_pred = model.predict(x_test.values)
    print("Modelo treinado. \n")
    
    # Resultados. (*100 para tratar como porcentagem)
    score = accuracy_score(y_test, y_pred)*100
    print("Acurácia do modelo: ", score)  # número de previsões corretas dividido pelo número total de previsões (fração de corretude).
    precision = precision_score(y_test, y_pred, pos_label='Benign')*100
    print("Precisão do modelo: ", precision)  # número de verdadeiros positivos dividido pelo número total de previsões positivas.
    recall = recall_score(y_test, y_pred, pos_label='Benign')*100
    print("Recall do modelo: ", recall) # número de verdadeiros positivos que foram encontrados.
    fscore = f1_score(y_test, y_pred, pos_label='Benign')*100
    print("F-Score do modelo: ", fscore) # média ponderada dos valores de precisão e recall.
    matrix = confusion_matrix(y_test, y_pred)
    print("\nMatriz de confusão: ")
    print(matrix)
    # Estruturação da matriz.
    # [  verdadeiro-negativo   falso-positivo ]
    # [ falso-negativo   verdadeiro-positivo  ]
    pickle.dump(model, open("./model_ml_fit.sav", 'wb')) # Salva em arquivo.

if __name__ == "__main__":
    treina_modelo(ajusta_dataset(df))